#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright 2012 Tobias Marschall
# 
# This file is part of CLEVER.
# 
# CLEVER is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# CLEVER is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with CLEVER.  If not, see <http://www.gnu.org/licenses/>.

from __future__ import print_function, division
from optparse import OptionParser
import sys
import os
from bitarray import bitarray
from collections import defaultdict
import math

__author__ = "Tobias Marschall"

usage = """%prog [options] <predictions> <truth> <chromosome-lengths>

Reads a file with predictions and a file with the true structural variations and
computes some statistics. To do that, a file with chromosome lengths is required
(format: "<chr-name> <length>", one chromosome per line).

The input files must contain one line per variation in the following format:
<chromosome> <start> <end> (INS|DEL) [(HOMO|HET)]
where <start> and <end> are 1-based and inclusive. In case of an insertion, 
<start> and <end> must have a difference of 1 and the insertion is assumed
to lie between these coordinates. The last (optional) column states whether
a variation is homo- or heterozygous."""

class VariationList:
	def __init__(self,filename):
		self.variations = defaultdict(list)
		n = 0
		for fields in (s.strip().split() for s in file(filename)):
			n += 1
			if (not 4 <= len(fields) <= 5) or (not fields[3] in ['INS','DEL']) or ((len(fields)==5) and (not fields[4] in ['HOMO','HET'])):
				print('Error parsing file "%s". Offending line: %d'%(filename,n), file=sys.stderr)
				sys.exit(1)
			chromosome, start, end, var_type = fields[0].lower(), int(fields[1]), int(fields[2]), fields[3]
			if not chromosome.startswith('chr'):
				chromosome = 'chr' + chromosome
			tags = set()
			# if information present, then add tag 'HOMO' or 'HET'
			if len(fields)==5:
				tags.add(fields[4])
			# convert from 1-based, inclusive coordinates to pythonic coordinates
			self.variations[chromosome].append((start-1, end, var_type, tags))
	def get(self, chromosome, index):
		"""Returns coordinates (coord1, coord2, var_type) for the variation with given index."""
		chromosome = chromosome.lower()
		coord1, coord2, var_type, homo = self.variations[chromosome][index]
		return (coord1, coord2, var_type)
	def get_deletion_bitarray(self, chromosome, chromosome_length, min_length=None, max_length=None):
		"""Returns a bitarray where all deleted positions are set to one."""
		chromosome = chromosome.lower()
		b = bitarray(chromosome_length)
		b.setall(0)
		for start, end, var_type, tags in self.variations[chromosome]:
			if var_type != 'DEL': continue
			length = end-start
			if (min_length != None) and (length < min_length): continue
			if (max_length != None) and (length > max_length): continue
			b[start:end] = True
		return b
	def find_all_deletion_overlaps(self, another_list, chromosome, min_length=None, max_length=None, tag=None):
		"""Returns a dictionary that maps each variation index of this list to a set of 
		indices of overlapping variations in the given list."""
		chromosome = chromosome.lower()
		events = []
		for n, (start, end, var_type, tags) in enumerate(self.variations[chromosome]):
			if var_type != 'DEL': continue
			if (tag!=None) and (not tag in tags): continue
			assert start < end
			length = end-start
			if (min_length != None) and (length < min_length): continue
			if (max_length != None) and (length > max_length): continue
			events.append((start, 'S', 0, n))
			events.append((end, 'E', 0, n))
		for n, (start, end, var_type, tags) in enumerate(another_list.variations[chromosome]):
			if var_type != 'DEL': continue
			assert start < end
			length = end-start
			if (min_length != None) and (length < min_length): continue
			if (max_length != None) and (length > max_length): continue
			events.append((start, 'S', 1, n))
			events.append((end, 'E', 1, n))
		events.sort()
		results = defaultdict(set)
		active = [set(), set()]
		for pos, event_type, list_idx, variation_index in events:
			if event_type == 'S':
				if list_idx == 0:
					for index1 in active[1]:
						results[variation_index].add(index1)
				else: 
					for index0 in active[0]:
						results[index0].add(variation_index)
				active[list_idx].add(variation_index)
			else:
				active[list_idx].remove(variation_index)
		return results
	def deletion_hit_statistics(self, another_list, chromosome, min_length=None, max_length=None, tag=None):
		"""Returns four lists, each length equals the number of deletions stored in this class (which match
		the given length constraints). 
		The first and second lists contain booleans saying whether the corresponding deletion is a "hit" or
		a "good hit", respectively. The third and the forth lists contain integers say how many of the 
		deletions given as "another_list" are covered by a "hit" or "good hit" so far."""
		chromosome = chromosome.lower()
		overlaps = self.find_all_deletion_overlaps(another_list, chromosome,min_length,max_length,tag)
		hit_list = []
		goodhit_list = []
		hit_count_list = []
		goodhit_count_list = []
		# set of variant indices that have already been covered by hits / good hits
		hit_set = set()
		goodhit_set = set()
		for i, (start0, end0, var_type0, tags0) in enumerate(self.variations[chromosome]):
			if var_type0 != 'DEL': continue
			if (tag!=None) and (not tag in tags0): continue
			length0 = end0 - start0
			if (min_length != None) and (length0 < min_length): continue
			if (max_length != None) and (length0 > max_length): continue
			is_hit = False
			is_good_hit = False
			for j in overlaps[i]:
				is_hit = True
				hit_set.add(j)
				start1, end1, var_type1, tags1 = another_list.variations[chromosome][j]
				length1 = end1 - start1	
				if (abs(start0-start1)<=20) and (abs(end0-end1)<=20) and (length0/length1<=1.1) and (length1/length0<=1.1):
					is_good_hit = True
					goodhit_set.add(j)
			hit_list.append(is_hit)
			goodhit_list.append(is_good_hit)
			hit_count_list.append(len(hit_set))
			goodhit_count_list.append(len(goodhit_set))
		return hit_list, goodhit_list, hit_count_list, goodhit_count_list
	def insertion_distances(self, another_list, chromosome, min_length=None, max_length=None, tag=None):
		"""Returns tree lists, "nearest_insertion", "own_length", and "partner_length", each containing one entry
		for each insertion stored in this class (for the given chromosome). For insertion with index i,
		nearest_insertion[i] gives the index of the nearest insertion in the other list and partner_length 
		gives its length."""
		chromosome = chromosome.lower()
		events = []
		n = 0
		lengths = []
		for start, end, var_type, tags in self.variations[chromosome]:
			# semantics for insertions...
			pos = start
			length = end
			if var_type != 'INS': continue
			if (tag!=None) and (not tag in tags): continue
			if length <= 0: continue
			if (min_length != None) and (length < min_length): continue
			if (max_length != None) and (length > max_length): continue
			events.append((pos, 0, n, length))
			lengths.append(length)
			n += 1
		insertion_count = n
		n = 0
		for start, end, var_type, tags in another_list.variations[chromosome]:
			# semantics for insertions...
			pos = start
			length = end
			if var_type != 'INS': continue
			if length <= 0: continue
			if (min_length != None) and (length < min_length): continue
			if (max_length != None) and (length > max_length): continue
			events.append((pos, 1, n, length))
			n += 1
		events.sort()
		partner_distances = [float('inf')] * insertion_count
		partner_lengths = [None] * insertion_count
		for i in xrange(2):
			last_pos = None
			last_length = None
			#print(events)
			for pos, event_type, idx, length in events:
				if event_type == 0:
					if last_pos != None:
						dist = abs(pos - last_pos)
						if dist < partner_distances[idx]:
							partner_distances[idx] = dist
							partner_lengths[idx] = last_length
				elif event_type == 1:
					last_pos = pos
					last_length = length
			events.reverse()
		return partner_distances, lengths, partner_lengths
	def insertion_hit_statistics(self, another_list, chromosome, min_length=None, max_length=None, tag=None):
		"""Returns two lists, each length equals the number of insertions stored in this class. 
		The lists contain booleans saying whether the corresponding deletion is a "hit" or
		a "good hit", respectively. """
		#print(zip(partner_distances,partner_lengths))
		chromosome = chromosome.lower()
		partner_distances, lengths, partner_lengths = self.insertion_distances(another_list, chromosome, min_length, max_length, tag)
		hits = []
		good_hits = []
		n = 0
		for dist, length in zip(partner_distances,partner_lengths):
			if dist <= 20:
				hits.append(True)
				good_hits.append((lengths[n]/partner_lengths[n] <= 1.1) and (partner_lengths[n]/lengths[n] <= 1.1))
			else:
				hits.append(False)
				good_hits.append(False)
			n += 1
		return hits, good_hits
	def add_tags(self, annotations, tag_name, read_length):
		"""Tags those variants for with the positions where reads to detect them are to be mapped
		are within  a distance of read_length or below to an annotation
		in the given AnnotationList with the given tag_name."""
		for chromosome, variation_list in self.variations.iteritems():
			if not annotations.annotations.has_key(chromosome): continue
			events = []
			for n, (coord1, coord2, var_type, tags) in enumerate(variation_list):
				if var_type == 'INS':
					if coord2 <= 0: continue
					events.append((coord1-read_length,'V0',n))
					events.append((coord1+read_length,'V1',n))
				elif var_type == 'DEL':
					if coord1 >= coord2: continue
					events.append((coord1-read_length,'V0',n))
					events.append((coord1,'V1',n))
					events.append((coord2,'V0',n))
					events.append((coord2+read_length,'V1',n))
			for n, (start,end) in enumerate(annotations.annotations[chromosome]):
				if start>end: continue
				events.append((start,'A0',n))
				events.append((end,'A1',n))
			events.sort()
			active_variants = set()
			active_annotations = set()
			for pos, event_type, index in events:
				if event_type == 'A0':
					for i in active_variants:
						variation_list[i][3].add(tag_name)
					active_annotations.add(index)
				elif event_type == 'A1':
					active_annotations.remove(index)
				elif event_type == 'V0':
					if len(active_annotations)>0:
						variation_list[index][3].add(tag_name)
					active_variants.add(index)
				elif event_type == 'V1':
					active_variants.remove(index)
			assert(len(active_variants)==0)
			assert(len(active_annotations)==0)

class AnnotationList:
	def __init__(self,filename):
		self.annotations = defaultdict(list)
		n = 0
		for fields in (s.strip().split() for s in file(filename)):
			n += 1
			if len(fields) != 3:
				print('Error parsing file "%s". Offending line: %d'(filename,n), file=sys.stderr)
				sys.exit(1)
			chromosome, start, end, = fields[0], int(fields[1]), int(fields[2])
			chromosome = chromosome.lower()
			# convert from 1-based, inclusive coordinates to pythonic coordinates
			self.annotations[chromosome].append((start-1, end))
	def get_bitarray(self, chromosome, chromosome_length):
		"""Returns a bitarray where all deleted positions are set to one."""
		chromosome = chromosome.lower()
		b = bitarray(chromosome_length)
		b.setall(0)
		for start, end in self.annotations[chromosome]:
			length = end-start
			b[start:end] = True
		return b

def format_stats(TP,FP,TN,FN):
	sensitivity = TP / (TP + FN) if (TP + FN)!=0 else float('nan')
	specificity = TN / (FP + TN) if (FP + TN)!=0 else float('nan')
	precision = TP / (TP + FP) if (TP + FP)!=0 else float('nan')
	# Matthew's correlation coefficient
	mcc_denom = math.sqrt((TP+FP)*(TP+FN)*(TN+FP)*(TN+FN))
	mcc = (TP*TN - FP*FN) / mcc_denom if mcc_denom!=0 else float('nan')
	return '%d %d %d %d %f %f %f %f'%(TP,FP,TN,FN,sensitivity,specificity,precision,mcc)

def scatter_hist(data, filename, title, x_label, y_label):
	"""Reads a list of pairs from "data" and plots a combined scatterplot/histogram to the given filename
	(in PDF) format."""
	# adapted from http://matplotlib.sourceforge.net/examples/pylab_examples/scatter_hist.html
	import numpy as np
	import matplotlib.pyplot as plt
	from matplotlib.ticker import NullFormatter
	if (data == None) or (len(data) == 0):
		plt.text(0.03, 0.8, "No data for \"%s\""%title)
		plt.savefig(filename, format='pdf') 
		plt.close()
		return          
	# the random data
	x = [xd for xd,yd in data]
	y = [yd for xd,yd in data]
	nullfmt   = NullFormatter()         # no labels
	# definitions for the axes
	left, width = 0.1, 0.65
	bottom, height = 0.1, 0.65
	bottom_h = left_h = left+width+0.02
	rect_scatter = [left, bottom, width, height]
	rect_histx = [left, bottom_h, width, 0.2]
	rect_histy = [left_h, bottom, 0.2, height]
	# start with a rectangular Figure
	plt.figure(1, figsize=(12,12))
	axScatter = plt.axes(rect_scatter)
	axHistx = plt.axes(rect_histx)
	axHisty = plt.axes(rect_histy)
	axHistx.xaxis.set_major_formatter(nullfmt)
	axHisty.yaxis.set_major_formatter(nullfmt)
	axScatter.scatter(x, y, marker='o', s=1)
	xmin = np.min(x)
	xmax = np.max(x)
	ymin = np.min(y)
	ymax = np.max(y)
	xymax = np.max( [np.max(np.fabs(x)), np.max(np.fabs(y))] )
	#xbinwidth = int((xmax-xmin)/70)
	#ybinwidth = int((ymax-ymin)/70)
	#lim = ( int(xymax/binwidth) + 1) * binwidth
	axScatter.set_xlim( (xmin-0.05*(xmax-xmin), xmax+0.05*(xmax-xmin)) )
	axScatter.set_ylim( (ymin-0.05*(ymax-ymin), ymax+0.05*(ymax-ymin)) )
	axScatter.set_xlabel(x_label)
	axScatter.set_ylabel(y_label)
	#xbins = np.arange(-xmin, xmax + xbinwidth, xbinwidth)
	#ybins = np.arange(-ymin, ymax + ybinwidth, ybinwidth)
	axHistx.hist(x, bins=70)
	axHisty.hist(y, bins=70, orientation='horizontal')
	axHistx.set_title(title)
	for tick in axHisty.xaxis.get_major_ticks():
		tick.label.set_rotation(270)
	axHistx.set_xlim( axScatter.get_xlim() )
	axHisty.set_ylim( axScatter.get_ylim() )
	plt.savefig(filename, format='pdf') 
	plt.close()

def print_deletion_stats(variants1, variants2, chromosomes, min_length, max_length, annotation, list_filename):
	total = 0 
	total_hits = 0
	total_good_hits = 0
	for chromosome in chromosomes:
		hit_list, goodhit_list, hit_count_list, goodhit_count_list = variants1.deletion_hit_statistics(variants2,chromosome,min_length,max_length, annotation)
		hits = sum(hit_list)
		good_hits = sum(goodhit_list)
		n = len(hit_list)
		total += n
		total_hits += hits
		total_good_hits += good_hits
		print(chromosome, n, hits, good_hits)
		if list_filename != None:
			out = open(list_filename,'w')
			hits = 0
			goodhits = 0
			i = 0
			for is_hit, is_good_hit, hit_count, goodhit_count in zip(hit_list, goodhit_list, hit_count_list, goodhit_count_list):
				hits += is_hit
				goodhits += is_good_hit
				i += 1
				print(is_hit, is_good_hit, hits, goodhits, hits/i, goodhits/i, hit_count, goodhit_count, hit_count/total_true_variants, goodhit_count/total_true_variants, file=out)
			out.close()
	print('total', total, total_hits, total_good_hits)

def print_insertion_stats(variants1, variants2, chromosomes, min_length, max_length, annotation):
	total = 0 
	total_hits = 0
	total_good_hits = 0
	for chromosome in chromosomes:
		hit_list, goodhit_list = variants1.insertion_hit_statistics(variants2, chromosome, min_length, max_length, annotation)
		hits = sum(hit_list)
		good_hits = sum(goodhit_list)
		n = len(hit_list)
		total += n
		total_hits += hits
		total_good_hits += good_hits
		print(chromosome, n, hits, good_hits)
	print('total', total, total_hits, total_good_hits)

def main():
	parser = OptionParser(usage=usage)
	parser.add_option("-c", action="store", dest="chromosomes", default=None,
					help="Comma separated list of chromosomes to consider (default=all).")
	parser.add_option("-m", action="store", dest="min_length", type=int, default=None,
					help="Minimum deletion length to be considered.")
	parser.add_option("-M", action="store", dest="max_length", type=int, default=None,
					help="Maximum deletion length to be considered.")
	parser.add_option("-a", action="store", dest="annotations", default=None,
					help="Comma separated list of annotation tracks. Separate performance statistics for all given tracks are printed. Tracks are assumed to be in format \"<chr> <start> <end>\", coordinates 1-based and inclusive.")
	parser.add_option("-l", action="store", dest="list_filename", default=None,
					help="Output a list of all deletion predictions matching the given length constraints to the given filename. Format \"<is_hit> <is_good_hit> <no_of_hits> <no_of_good_hits>\".")
	parser.add_option("-d", action="store", dest="deletion_plots", default=None,
					help="Plot accuracy of deletion predictions (in PDF format) to the given filename.")
	parser.add_option("-i", action="store", dest="insertion_plots", default=None,
					help="Plot accuracy of insertion predictions (in PDF format) to the given filename.")

	(options, args) = parser.parse_args()
	if (len(args)!=3):
		parser.print_help()
		sys.exit(1)
	annotations = {}
	if options.annotations != None:
		for filename in options.annotations.split(','):
			name = os.path.basename(filename)
			annotations[name] = AnnotationList(filename)
	if (options.deletion_plots != None) or (options.insertion_plots != None):
		import matplotlib
		matplotlib.use('pdf')
	predictions = VariationList(args[0])
	true_variants = VariationList(args[1])
	for name, annotation_list in annotations.iteritems():
		# TODO: Don't use hard-coded distance here
		predictions.add_tags(annotation_list, name, 312 - 100)
		true_variants.add_tags(annotation_list, name, 312 - 100)
	chromosome_lengths = dict([(f[0],int(f[1])) for f in (s.split() for s in open(args[2]))])
	if options.chromosomes != None:
		chromosomes = set(options.chromosomes.split(','))
	else:
		chromosomes = set(chromosome_lengths.keys())
	if (options.list_filename != None) and (len(chromosomes)>1):
		print("Error: option -l can only be used for a single chromosome, use together with option -c", file=sys.stderr)
		sys.exit(1)
	for annotation in [None] + annotations.keys():
		print('=====================================================================')
		print('TRUE-DELETION-WISE STATISTICS (%s)'%annotation)
		print_deletion_stats(true_variants, predictions, chromosomes, options.min_length, options.max_length, annotation, None)
		print('\n=====================================================================')
		print('PREDICTED-DELETION-WISE STATISTICS (%s)'%annotation)
		print_deletion_stats(predictions, true_variants, chromosomes, options.min_length, options.max_length, annotation, options.list_filename)
		print('\n=====================================================================')
		print('TRUE-INSERTION-WISE STATISTICS (%s)'%annotation)
		print_insertion_stats(true_variants, predictions, chromosomes, options.min_length, options.max_length, annotation)
		print('\n=====================================================================')
		print('PREDICTED-INSERTION-WISE STATISTICS (%s)'%annotation)
		print_insertion_stats(predictions, true_variants, chromosomes, options.min_length, options.max_length, annotation)
	print('\n=====================================================================')
	#print('POSITION-WISE DELETION STATISTICS')
	#for annotation, invert_annotation in [(None,False)]+[(key,b) for key in annotations.keys() for b in [False,True]]:
		#print('---------------------------------------------------------------------')
		#print('Annotation: %s'%('NOT' if invert_annotation else ''), annotation)
		#total_TP = 0
		#total_FP = 0
		#total_TN = 0
		#total_FN = 0
		#for chromosome in chromosomes:
			#predition_bitarray = predictions.get_deletion_bitarray(chromosome,chromosome_lengths[chromosome],options.min_length,options.max_length)
			#truth_bitarray = true_variants.get_deletion_bitarray(chromosome,chromosome_lengths[chromosome],options.min_length,options.max_length)
			#if annotation == None:
				#TP = ( predition_bitarray &  truth_bitarray).count()
				#FP = ( predition_bitarray & ~truth_bitarray).count()
				#TN = (~predition_bitarray & ~truth_bitarray).count()
				#FN = (~predition_bitarray &  truth_bitarray).count()
			#else:
				#annotation_bitarray = annotations[annotation].get_bitarray(chromosome,chromosome_lengths[chromosome])
				#if invert_annotation:
					#annotation_bitarray = ~annotation_bitarray
				#TP = ( predition_bitarray &  truth_bitarray & annotation_bitarray).count()
				#FP = ( predition_bitarray & ~truth_bitarray & annotation_bitarray).count()
				#TN = (~predition_bitarray & ~truth_bitarray & annotation_bitarray).count()
				#FN = (~predition_bitarray &  truth_bitarray & annotation_bitarray).count()
			#total_TP += TP
			#total_FP += FP
			#total_TN += TN
			#total_FN += FN
			#print(chromosome, 'pos-wise:', format_stats(TP,FP,TN,FN))
		#print('total', 'pos-wise:', format_stats(total_TP,total_FP,total_TN,total_FN))
		#print('\n=====================================================================')
	
	# create plots of deletion predition accuracy?
	if options.deletion_plots != None:
		# for every hitting, i.e. overlapping, deletion, store a pair (offset start, offset end), given
		# how far predicted start and end are away from true start and end.
		deletion_accuracy = []
		for chromosome in chromosomes:
			overlaps = predictions.find_all_deletion_overlaps(true_variants, chromosome, options.min_length, options.max_length)
			for n, overlapping_set in overlaps.iteritems():
				if len(overlapping_set) == 0: continue
				#print(n, '-->', overlapping_set)
				diff_pairs = []
				pred_start, pred_end, pred_type = predictions.get(chromosome,n)
				for x in overlapping_set:
					true_start, true_end, true_type = true_variants.get(chromosome,x)
					assert(pred_type == true_type == "DEL")
					diff_pairs.append((pred_start-true_start,pred_end-true_end))
				if len(diff_pairs) > 1:
					# sort by sum of differences, i.e. if ambiguous, take the better one
					diff_pairs.sort(cmp=lambda x,y: cmp(abs(x[0])+abs(x[1]),abs(y[0])+abs(y[1])))
				deletion_accuracy.append(diff_pairs[0])
		scatter_hist(deletion_accuracy,options.deletion_plots, "Deletions %s (length %d-%d)"%(args[0],options.min_length,options.max_length), "Start coordinate difference", "End coordinate difference")

	# create plots of deletion predition accuracy?
	if options.insertion_plots != None:
		# for every insertion, store a pair (distance, length diff) with information on the positions
		# of the nearest true insertion
		insertion_accuracy = []
		for chromosome in chromosomes:
			partner_distances, lengths, partner_lengths = predictions.insertion_distances(true_variants, chromosome, options.min_length, options.max_length)
			for n in xrange(len(partner_distances)):
				pos_diff = min(partner_distances[n],250)
				length_diff = min(abs(lengths[n]-partner_lengths[n]),250) if partner_lengths[n]!=None else 250
				insertion_accuracy.append((pos_diff,length_diff))
		scatter_hist(insertion_accuracy,options.insertion_plots, "Insertions %s (length %d-%d)"%(args[0],options.min_length,options.max_length), "Position difference", "Length difference")

if __name__ == '__main__':
	sys.exit(main())
